#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created by PyCharm.

@Date    : Thu Feb 25 2021 
@Time    : 09:17:17
@File    : train.py
@Author  : alpha
"""


import torch
import albumentations as AT

from torch import optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader

from torchvision.models import resnet18

from src.log import logger
from src.dataset import CelebASpoofCroppedDataset
from src.model import FaceSpoofNet
from src.regvgg import RepVGG18, RepVGG18Slim
from src.resnet import ResNet18
from src.trainer import train_feat, train_cls, train_join
from src.losses import MetricLoss, TripleBCELoss


if __name__ == '__main__':

    transform = AT.Compose([
        AT.HorizontalFlip(p=0.5),
        AT.MotionBlur(p=0.2),
        AT.Rotate(limit=20, p=0.2),
        AT.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.2),
        AT.RandomResizedCrop(256, 256, scale=(0.8, 1.0), ratio=(1, 1), p=0.2)
    ])

    train_dataset = CelebASpoofCroppedDataset(
        root='/data/dataset/CelebA_Spoof_cropped', mode='train', transform=transform
    )
    test_dataset = CelebASpoofCroppedDataset(
        root='/data/dataset/CelebA_Spoof_cropped', mode='test'
    )

    train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
    test_dataloader = DataLoader(test_dataset, batch_size=32, num_workers=2)

    # TODO: Implement MultiGPU training
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = FaceSpoofNet(feat_channels=512, backbone=ResNet18)

    # JOINED TRAINING
    # model.load_state_dict(torch.load("checkpoints/face_spoof_with_cls_epoch000.pth"))
    #
    # met_criterion = MetricLoss()
    # cls_criterion = TripleBCELoss()
    # # optimizer = optim.Adam(params=model.parameters(), lr=0.001)
    # optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.00001)
    # scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[5, 10, 15], gamma=0.1)
    #
    # logger.info('Joined Training Starts...')
    # train_join(
    #     model,
    #     device,
    #     met_criterion,
    #     cls_criterion,
    #     optimizer,
    #     train_dataloader,
    #     test_dataloader,
    #     num_epoches=20,
    #     scheduler=scheduler
    # )
    # logger.info('Joined Training Finished!!!')

    # load pretrain model
    # pretrain_model = 'checkpoints/repvgg18_20210227_v2/face_spoof_with_cls_best.pth'
    # pretrain_model = 'checkpoints/repvgg18_20210227_v2_finetune/face_spoof_epoch003_wo_cls.pth'
    # pretrain_model = 'checkpoints/repvgg18_20210227_v2_finetune/face_spoof_with_cls_epoch001.pth'
    # logger.info('Load pretrain model: {}'.format(pretrain_model))
    # model.load_state_dict(torch.load(pretrain_model), strict=False)
    # model.load_state_dict((torch.load('checkpoints/repvgg18/face_spoof_epoch006_wo_cls.pth')))
    # model.load_state_dict((torch.load('checkpoints/repvgg18_20210227/face_spoof_epoch008_wo_cls.pth')))

    # STAGE ONE
    logger.info('STAGE ONE Training Starts...')
    criterion = MetricLoss()
    optimizer = optim.AdamW(params=model.parameters(), lr=0.01)
    # optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9, weight_decay=0.00001)
    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[10, 13, 14], gamma=0.1)
    scheduler = None
    train_feat(
        model,
        device,
        criterion,
        optimizer,
        train_dataloader,
        num_epoches=15,
        scheduler=scheduler,
        save_path='checkpoints/resnet18',
        fixbn=True
    )
    logger.info('STAGE ONE Training Finished!!!')


    # STAGE TWO
    logger.info('STAGE TWO Training Starts...')
    criterion = TripleBCELoss()
    optimizer = optim.AdamW(params=model.parameters(), lr=0.01)
    # optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9, weight_decay=0.00001)
    # scheduler = None
    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[3, 4], gamma=0.1)
    train_cls(
        model,
        device,
        criterion,
        optimizer,
        train_dataloader,
        test_dataloader,
        num_epoches=5,
        scheduler=scheduler,
        save_path='checkpoints/resnet18',
        fixbn=True
    )
    logger.info('STAGE TWO Training Finished!!!')
